/*
 * Copyright (c) 2011-2018, Meituan Dianping. All Rights Reserved.
 *
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.dianping.zebra.shard.jdbc.parallel;

import com.dianping.zebra.group.util.DaoContextHolder;

import java.sql.SQLException;
import java.sql.Statement;
import java.util.Map;
import java.util.concurrent.Callable;

public class StatementExecuteUpdateCallable implements Callable<UpdateResult> {

	private Map<String, Statement> sqlStmtMap;

	private int autoGeneratedKeys;

	private int[] columnIndexes;

	private String[] columnNames;

	protected TaskExecuteResult taskExecuteResult;

	public StatementExecuteUpdateCallable(Map<String, Statement> sqlStmtMap, int autoGeneratedKeys, int[] columnIndexes,
	      String[] columnNames) {
		super();
		this.sqlStmtMap = sqlStmtMap;
		this.autoGeneratedKeys = autoGeneratedKeys;
		this.columnIndexes = columnIndexes;
		this.columnNames = columnNames;
	}

	@Override
	public UpdateResult call() throws Exception {
		return new UpdateResult(stmtExecuteUpdate());
	}

	private int stmtExecuteUpdate() throws SQLException {
		try {
			int affectedRows = 0;

			for (Map.Entry<String, Statement> entry : sqlStmtMap.entrySet()) {
				String sql = entry.getKey();
				Statement stmt = entry.getValue();

				if (autoGeneratedKeys == -1 && columnIndexes == null && columnNames == null) {
					affectedRows += stmt.executeUpdate(sql, Statement.RETURN_GENERATED_KEYS);
				} else if (autoGeneratedKeys != -1) {
					affectedRows += stmt.executeUpdate(sql, autoGeneratedKeys);
				} else if (columnIndexes != null) {
					affectedRows += stmt.executeUpdate(sql, columnIndexes);
				} else if (columnNames != null) {
					affectedRows += stmt.executeUpdate(sql, columnNames);
				} else {
					affectedRows += stmt.executeUpdate(sql, Statement.RETURN_GENERATED_KEYS);
				}
			}

			return affectedRows;
		} finally {
			DaoContextHolder.clearSqlName();
		}
	}

	public TaskExecuteResult getTaskExecuteResult() {
		return taskExecuteResult;
	}

	public void setTaskExecuteResult(TaskExecuteResult taskExecuteResult) {
		this.taskExecuteResult = taskExecuteResult;
	}
}
